# Name: Xueyan Lyu
import numpy as np
import pandas as pd
import datetime
from sklearn import svm

# set the model
# I chose SVM regression model
model_svm = svm.SVR()
# read csv file
pr = pd.read_csv("owid-covid-data.csv")

# select data we need
US_data = []
start_date = datetime.datetime.strptime("2021-06-25", '%Y-%m-%d')
end_date = datetime.datetime.strptime("2021-08-10", '%Y-%m-%d')
for i in range(len(pr)):
    line = pr.iloc[i]
    if line['location'] == 'United States':
        date = line['date']
        if start_date <= datetime.datetime.strptime(date,'%Y-%m-%d') <= end_date:
            US_data.append(line['new_cases'])

model = model_svm
x_train = np.linspace(0, len(US_data)-1, len(US_data))
# prepare the data
x_train = np.array(x_train).reshape(-1,1)
y_train = np.array(US_data).reshape(-1,1)
# fit the model
model.fit(x_train, y_train)
# suppose 2021.6.25 is day 0, then calculate 2021.9.1
pre_date = datetime.datetime.strptime("2021-09-01", '%Y-%m-%d')
start_date = datetime.datetime.strptime("2021-06-25", '%Y-%m-%d')
expand = (pre_date-start_date).days
x_test = np.array([expand]).reshape(-1, 1)
# get the prediction
result = model.predict(x_test)
print(result)
# result = [29892.77263722]